# -*- coding: utf-8 -*-
# @Time : 2020/10/31 19:51
# @Author : wudeyang
# @email :wudeyang@sjtu.edu.cn
# @Description:

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class Maploss(nn.Module):
    def __init__(self,ohem_ratio=3):
        self.ohem_ratio=ohem_ratio

        super(Maploss,self).__init__()


    def ohem_single(self, score, gt_mask):
        """
        # 此处参考了pan的做法
        :param score: 预测的高斯结果
        :param gt_text: gt的高斯结果
        :param
        :return:
        """
        pos_num = (int)(np.sum(gt_mask > 0.5)) # 正样本点的数量

        if pos_num == 0:
            selected_mask = gt_mask
            selected_mask = selected_mask.reshape(1, selected_mask.shape[0], selected_mask.shape[1]).astype('float32')
            return selected_mask

        neg_num = (int)(np.sum(gt_mask <= 0.5)) # 不是文本的点的数量
        neg_num = (int)(min(pos_num * self.ohem_ratio, neg_num)) # 控制负样本点的数量和正样本点的数量是3：1

        if neg_num == 0:
            selected_mask = gt_mask
            selected_mask = selected_mask.reshape(1, selected_mask.shape[0], selected_mask.shape[1]).astype('float32')
            return selected_mask

        neg_score = score[gt_mask <= 0.5] # 过滤score 图，过滤条件：没有文本,生成一行数据
        neg_score_sorted = np.sort(-neg_score)
        threshold = -neg_score_sorted[neg_num - 1]# 从大到小找出来 di neg_num-1个点当作阈值
        selected_mask = (score >= threshold) | (gt_mask > 0.5) # 选出来负样本和全部的正样本
        selected_mask = selected_mask.reshape(1, selected_mask.shape[0], selected_mask.shape[1]).astype('float32')
        return selected_mask

    def ohem_batch(self, scores, gt_masks):
        scores = scores.data.cpu().numpy()
        gt_texts = gt_masks.data.cpu().numpy()


        selected_masks = []
        for i in range(scores.shape[0]):
            selected_masks.append(self.ohem_single(scores[i, :, :], gt_texts[i, :, :]))

        selected_masks = np.concatenate(selected_masks, 0)
        selected_masks = torch.from_numpy(selected_masks).float()

        return selected_masks


    def dice_loss(self, input, target, mask):
        """

        :param input: 预测的text
        :param target: gt_text
        :param mask: maks图，经过ohem
        :return:
        """
        input = torch.sigmoid(input)
        # 生成二值化图
        target[target <= 0.5] = 0
        target[target > 0.5] = 1
        input = input.contiguous().view(input.size()[0], -1)
        target = target.contiguous().view(target.size()[0], -1)
        mask = mask.contiguous().view(mask.size()[0], -1)

        input = input * mask
        target = target * mask

        a = torch.sum(input * target, 1)
        b = torch.sum(input * input, 1) + 0.001
        c = torch.sum(target * target, 1) + 0.001
        d = (2 * a) / (b + c)
        return 1 - d



    def forward(self, GT_score, pre_score,GT_mask):
        """
        :param GT_score: GT的高斯图 B,H,W
        :param pre_score: 网络输出的高斯分布图
        :param GT_mask: 原图的GT图
        :return:
        """
        # 先进行生维度,再进行插值
        texts=pre_score[0].permute(3,0,1,2)
        texts=F.interpolate(texts, size=list(np.array(texts.size()[2:])*2), mode="bilinear")
        texts=texts[0,:,:,:]
        selected_masks = self.ohem_batch(texts, GT_mask)
        selected_masks = selected_masks.cuda()
        loss_fn = torch.nn.MSELoss()
        GT_score=selected_masks*GT_score
        texts=selected_masks*texts
        loss=loss_fn(GT_score,texts)
        return loss


if __name__=='__main__':
#
#     preds[0]torch.Size([1, 225, 300, 1])
#     images, masks, gauss:etorch.Size([1, 3, 450, 600]) torch.Size([1, 450, 600]) torch.Size([1, 450, 600])

    pre=torch.tensor([[0.9,0.7],[0.6,0.4]],dtype=torch.float32)
    gt=torch.ones((1,4,4),dtype=torch.float32)
    mask=torch.tensor([[1,1,0,1],[0,0,0,0],[0,0,0,0],[0,0,0,0]],dtype=torch.float32)
    # 升维度
    print(pre)
    print(gt)
    print(mask)
    pre=pre.unsqueeze(0)
    pre=pre.unsqueeze(3)
    mask=mask.unsqueeze(0)
    loss=Maploss()
    print('*'*10)
    l=loss.forward(gt,[pre,0],mask)
    print(pre)
    print(gt)
    print(mask)
    print(l)
